library("AER")
library("mc2d")

set.seed(123)

simIV <- function(n = 200,
	mu = 0,
	kappa = 1, # treatment strength
	lambda = 0, # responsiveness
	hetero = 1,
	cov = 0.5,
	sigma = 2 # std for errors
	) {
	# errors
	errors <- MASS::mvrnorm(n, c(1, 1), sigma^2* matrix(c(1, cov, cov, 1), 2, 2))
	coefs <- MASS::mvrnorm(n, c(2, 1), hetero * 0.25* sigma^2* matrix(c(1, lambda, lambda, 0.5), 2, 2))
	v <- errors[, 1]
	u <- errors[, 2]
	# coefficients
	beta <- coefs[,1]
	pi <- coefs[,2]
	# instrument, treatment, outcome
	z <- mc2d::rbern(n, 0.5)
	delta <- sapply(sapply(kappa * pi, min, 1), max, 0)
	x <- ifelse(delta * z + (1-delta) * rnorm(n, 0, 1) + 0.2*v>0, 1, 0)
	y <- 5 + beta*x + mu * z + u + rnorm(n, 0, 1)
	# save
	out <- cbind.data.frame(y, x, z, beta, pi, delta, u, v)
	return(out)
}

runSim <- function(nsims = 10, mu = 0, lambdas, kappas, hetero = hetero, cov = 0.5) {
	results <- matrix(NA, length(lambdas) * length(kappas) * nsims, 7) # save coefs, ratio and rho
	colnames(results) <- c("Coef.OLS","Coef.2SLS","ratio","ratio2","rho","sig.OLS","sig.IV")
	p <- 1
	for (lambda in lambdas) {
		for (kappa in kappas) {
			for (k in 1:nsims) {
				d <- simIV(n = 1000, mu = mu, lambda = lambda, kappa = kappa, hetero = hetero, cov = cov)
				# OLS
				out.ols <- lm(y ~ x, data = d)
				coef.ols <- out.ols$coefficients[2]
				pvalue <- summary(out.ols)$coefficients[2,4]
				sig.ols <- ifelse(pvalue<0.05, 1, 0)
				# IV
				out.iv <- ivreg(y ~ x | z, data = d)
				coef.iv <- out.iv$coefficients[2]
				pvalue <- summary(out.iv)$coefficients[2,4]
				sig.iv <- ifelse(pvalue<0.05, 1, 0)
				# result
				ratio <- abs(coef.iv/coef.ols)
				ratio2 <- abs(coef.iv - coef.ols)/abs(coef.ols)
				rho <- cor(d$x, d$z)
				results[p,] <- c(coef.ols, coef.iv, ratio, ratio2, rho, sig.ols, sig.iv)
				p <- p + 1
				if (p%%100==0) {cat(".")}
			}
		}
	}
	return(results)
}


#########################################################
## No Heterogeneity (Positive Selection)
#########################################################

set.seed(123)
out.b1 <- runSim(nsims = 100, mu = 0, lambda = 0.7, hetero = 0, cov = 0.7,
	kappas = seq(0.25, 1, length.out = 20))

# Exclusion restriction failure
out.b2 <- runSim(nsims = 100, mu = 1, lambda = 0.7, hetero = 0, cov = 0.7,
	kappas = seq(0.25, 1, length.out = 20))

# HTE
out.b3 <- runSim(nsims = 100, mu = 0, lambda = 0.7, hetero = 0.3, cov = 0.7,
	kappas = seq(0.25, 1, length.out = 20))

cat("Done\n")


save(out.b1, out.b2, out.b3, file = "graphs/sim_binary.RData")



#########################################################
## Graphics
#########################################################

load("graphs/sim_binary.RData")
library("latex2exp")
xlab <- TeX(r'($|\hat{\rho}(d, \hat{d})|$)')
ylab <- TeX(r'($|\hat{\tau}_{2SLS} / \hat{\tau}_{OLS}|$)')

# benchmark
coef.ols <- out.b1[,1]
coef.iv <- out.b1[,2]
ratio <-out.b1[,3]
ratio2 <-out.b1[,4]
rho <- out.b1[,5]
sig <- ifelse(out.b1[,7]==1, TRUE, FALSE)

pdf("graphs/FigureA6c.pdf")
par(mar = c(4,4,1,1))
plot(1, type = "n", xlim = c(0.0, 0.5), ylim = c(-1.5,1.5), axes = FALSE,
		xlab = "", ylab = "")
points(abs(rho), log10(abs(ratio)), col = "#000000", cex = 0.8)
reg <- lm(log10(abs(ratio))~abs(rho))
abline(reg, lwd = 2, col = 2)
y0 <- c(0.001, 0.01, 0.03, 0.1, 0.3, 1, 3, 10, 30, 100)
axis(1); box()
axis(2, at = log10(y0), labels = y0)
f <- summary(reg)$fstatistic
p <- pf(f[1],f[2],f[3],lower.tail=FALSE)
text(0.45, 1.45, paste("R2:",sprintf("%.3f",summary(reg)$r.squared)), cex = 1.2)
mtext(xlab, 1, 3)
mtext(ylab, 2, 2.5)
graphics.off()

# publication bias
pdf("graphs/FigureA7b.pdf")
par(mar = c(4,4,1,1))
plot(1, type = "n", xlim = c(0.0, 0.5), ylim = c(-1.5,1.5), axes = FALSE,
		xlab = "", ylab = "")
points(abs(rho[sig]), log10(abs(ratio[sig])), col = "#000000", cex = 0.8)
reg <- lm(log10(abs(ratio[sig]))~abs(rho[sig]))
abline(reg, lwd = 2, col = 2)
y0 <- c(0.001, 0.01, 0.03, 0.1, 0.3, 1, 3, 10, 30, 100)
axis(1); box()
axis(2, at = log10(y0), labels = y0)
f <- summary(reg)$fstatistic
p <- pf(f[1],f[2],f[3],lower.tail=FALSE)
text(0.45, 1.45, paste("R2:",sprintf("%.3f",summary(reg)$r.squared)), cex = 1.2)
mtext(xlab, 1, 3)
mtext(ylab, 2, 2.5)
graphics.off()

# exclusion restriction failure
coef.ols <- out.b2[,1]
coef.iv <- out.b2[,2]
ratio <-out.b2[,3]
ratio2 <-out.b2[,4]
rho <- out.b2[,5]
sig <- ifelse(out.b2[,7]==1, TRUE, FALSE)

pdf("graphs/FigureA6d.pdf")
par(mar = c(4,4,1,1))
plot(1, type = "n", xlim = c(0.0, 0.5), ylim = c(-1.5,1.5), axes = FALSE,
		xlab = "", ylab = "")
points(abs(rho), log10(abs(ratio)), col = "#000000", cex = 0.8)
reg <- lm(log10(abs(ratio))~abs(rho))
abline(reg, lwd = 2, col = 2)
y0 <- c(0.001, 0.01, 0.03, 0.1, 0.3, 1, 3, 10, 30, 100)
axis(1); box()
axis(2, at = log10(y0), labels = y0)
f <- summary(reg)$fstatistic
p <- pf(f[1],f[2],f[3],lower.tail=FALSE)
text(0.45, 1.45, paste("R2:",sprintf("%.3f",summary(reg)$r.squared)), cex = 1.2)
mtext(xlab, 1, 3)
mtext(ylab, 2, 2.5)
graphics.off()

# HTE
coef.ols <- out.b3[,1]
coef.iv <- out.b3[,2]
ratio <-out.b3[,3]
ratio2 <-out.b3[,4]
rho <- out.b3[,5]
sig <- ifelse(out.b3[,7]==1, TRUE, FALSE)

pdf("graphs/FigureA8c.pdf")
par(mar = c(4,4,1,1))
plot(1, type = "n", xlim = c(0, 0.5), ylim = c(-2.5,1.5), axes = FALSE,
		xlab = "", ylab = "")
points(abs(rho), log10(abs(ratio)), col = "#000000", cex = 0.8)
reg <- lm(log10(abs(ratio))~abs(rho))
abline(reg, lwd = 2, col = 2)
y0 <- c(0.001, 0.01, 0.03, 0.1, 0.3, 1, 3, 10, 30, 100)
axis(1); box()
axis(2, at = log10(y0), labels = y0)
f <- summary(reg)$fstatistic
p <- pf(f[1],f[2],f[3],lower.tail=FALSE)
text(0.45, 1.45, paste("R2:",sprintf("%.3f",summary(reg)$r.squared)), cex = 1.2)
mtext(xlab, 1, 3)
mtext(ylab, 2, 2.5)
graphics.off()

pdf("graphs/FigureA8d.pdf")
par(mar = c(4,4,1,1))
plot(1, type = "n", xlim = c(0, 0.5), ylim = c(-2.5,1.5), axes = FALSE,
		xlab = "", ylab = "")
points(abs(rho[sig]), log10(abs(ratio[sig])), col = "#000000", cex = 0.8)
reg <- lm(log10(abs(ratio[sig]))~abs(rho[sig]))
abline(reg, lwd = 2, col = 2)
y0 <- c(0.001, 0.01, 0.03, 0.1, 0.3, 1, 3, 10, 30, 100)
axis(1); box()
axis(2, at = log10(y0), labels = y0)
f <- summary(reg)$fstatistic
p <- pf(f[1],f[2],f[3],lower.tail=FALSE)
text(0.45, 1.45, paste("R2:",sprintf("%.3f",summary(reg)$r.squared)), cex = 1.2)
mtext(xlab, 1, 3)
mtext(ylab, 2, 2.5)
graphics.off()


